{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "MHA、GQA、MQA、MLA\n",
    "\n",
    "DeepSeek-LLM模型中用的是GQA，也就是grouped-query attention，GQA表示分组查询注意力机制。\n",
    "\n",
    "DeepSeek V2、V3用的是MLA，也就是Multi-head Latent Attention，也就是多头潜在注意力机制。\n",
    "DeepSeek R1是基于V3通过强化学习训练出来的，所以也是用的MLA。\n",
    "\n",
    "GQA论文：https://arxiv.org/abs/2305.13245"
   ],
   "id": "75ca5623955f57f2"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.430341Z",
     "start_time": "2025-07-26T00:51:28.890639Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "\n",
    "# x表示一个输入的seq，有4个词，每个词向量长度为8\n",
    "x = torch.rand(4, 8)\n",
    "\n",
    "# w是二维矩阵\n",
    "wq_head1 = torch.rand(8, 8)\n",
    "wq_head2 = torch.rand(8, 8)\n",
    "wq_head3 = torch.rand(8, 8)\n",
    "wq_head4 = torch.rand(8, 8)\n",
    "\n",
    "wk_group1 = torch.rand(8, 8)\n",
    "wk_group2 = torch.rand(8, 8)\n",
    "\n",
    "wv_group1 = torch.rand(8, 8)\n",
    "wv_group2 = torch.rand(8, 8)\n",
    "\n",
    "q_head1 = torch.matmul(x, wq_head1)\n",
    "q_head2 = torch.matmul(x, wq_head2)\n",
    "q_head3 = torch.matmul(x, wq_head3)\n",
    "q_head4 = torch.matmul(x, wq_head4)\n",
    "\n",
    "k_group1 = torch.matmul(x, wk_group1) # 第1个Q和第2个Q对应第1组，第3个Q和第4个Q对应第2组\n",
    "k_group2 = torch.matmul(x, wk_group2)\n",
    "\n",
    "v_group1 = torch.matmul(x, wv_group1)\n",
    "v_group2 = torch.matmul(x, wv_group2)"
   ],
   "id": "7ee18446f5ae7452",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.485546Z",
     "start_time": "2025-07-26T00:51:29.482742Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk_score_head1 = torch.matmul(q_head1, k_group1.T)\n",
    "dk_head1 = k_group1.size(-1)\n",
    "qk_weight_head1 = torch.softmax(qk_score_head1 / torch.sqrt(torch.tensor(dk_head1)), dim=-1)\n",
    "o_head1 = torch.matmul(qk_weight_head1, v_group1)"
   ],
   "id": "7b7d504b707c1b3",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.499044Z",
     "start_time": "2025-07-26T00:51:29.496965Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk_score_head2 = torch.matmul(q_head2, k_group1.T)\n",
    "dk_head2 = k_group1.size(-1)\n",
    "qk_weight_head2 = torch.softmax(qk_score_head2 / torch.sqrt(torch.tensor(dk_head2)), dim=-1)\n",
    "o_head2 = torch.matmul(qk_weight_head2, v_group1)"
   ],
   "id": "7eebecc244dbf577",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.512112Z",
     "start_time": "2025-07-26T00:51:29.509686Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk_score_head3 = torch.matmul(q_head3, k_group2.T)\n",
    "dk_head3 = k_group2.size(-1)\n",
    "qk_weight_head3 = torch.softmax(qk_score_head3 / torch.sqrt(torch.tensor(dk_head3)), dim=-1)\n",
    "o_head3 = torch.matmul(qk_weight_head3, v_group2)"
   ],
   "id": "d323b45347ea125d",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.526316Z",
     "start_time": "2025-07-26T00:51:29.523948Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk_score_head4 = torch.matmul(q_head4, k_group2.T)\n",
    "dk_head4 = k_group2.size(-1)\n",
    "qk_weight_head4 = torch.softmax(qk_score_head4 / torch.sqrt(torch.tensor(dk_head4)), dim=-1)\n",
    "o_head4 = torch.matmul(qk_weight_head4, v_group2)"
   ],
   "id": "89ec68f4429bc6d5",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.542655Z",
     "start_time": "2025-07-26T00:51:29.538378Z"
    }
   },
   "cell_type": "code",
   "source": [
    "o = torch.cat([o_head1, o_head2, o_head3, o_head4], dim=1)\n",
    "o"
   ],
   "id": "38e503043ee37c12",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.4892, 2.2409, 3.0830, 1.9829, 1.3328, 1.7617, 2.1029, 2.2179, 2.5397,\n",
       "         2.2371, 3.0668, 1.9731, 1.3336, 1.7150, 2.1268, 2.2563, 1.7857, 2.2017,\n",
       "         2.5348, 1.9395, 2.0288, 2.7907, 2.6411, 2.9243, 1.7816, 2.1895, 2.4880,\n",
       "         1.8965, 2.0238, 2.7434, 2.6256, 2.9122],\n",
       "        [2.5552, 2.2504, 3.1093, 1.9817, 1.3569, 1.7436, 2.1377, 2.2815, 2.5890,\n",
       "         2.2504, 3.1021, 1.9783, 1.3595, 1.7155, 2.1530, 2.3079, 1.7906, 2.2065,\n",
       "         2.5828, 1.9795, 2.0201, 2.8361, 2.6535, 2.9298, 1.7834, 2.2026, 2.5204,\n",
       "         1.9290, 2.0361, 2.7784, 2.6378, 2.9247],\n",
       "        [2.5334, 2.2508, 3.1048, 1.9868, 1.3515, 1.7533, 2.1251, 2.2612, 2.6769,\n",
       "         2.2639, 3.1331, 1.9789, 1.3901, 1.6876, 2.1978, 2.3910, 1.7895, 2.2071,\n",
       "         2.5700, 1.9696, 2.0252, 2.8246, 2.6513, 2.9300, 1.7861, 2.2060, 2.5471,\n",
       "         1.9517, 2.0326, 2.8040, 2.6451, 2.9284],\n",
       "        [2.4949, 2.2395, 3.0766, 1.9814, 1.3308, 1.7523, 2.1050, 2.2209, 2.5820,\n",
       "         2.2445, 3.0913, 1.9724, 1.3524, 1.7103, 2.1503, 2.2993, 1.7857, 2.1968,\n",
       "         2.5180, 1.9231, 2.0244, 2.7729, 2.6365, 2.9196, 1.7820, 2.1948, 2.4871,\n",
       "         1.8982, 2.0324, 2.7443, 2.6286, 2.9171]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-26T00:51:29.579718Z",
     "start_time": "2025-07-26T00:51:29.576369Z"
    }
   },
   "cell_type": "code",
   "source": [
    "wo = torch.rand(4*8, 8)\n",
    "out_put = torch.matmul(o, wo)\n",
    "out_put"
   ],
   "id": "18cc2965ef964b2",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[35.3237, 32.7969, 37.6285, 39.0346, 41.6376, 33.9676, 38.3575, 40.0775],\n",
       "        [35.6853, 33.1366, 38.0081, 39.4083, 42.0654, 34.2874, 38.7489, 40.4958],\n",
       "        [35.7997, 33.3191, 38.1082, 39.5443, 42.1714, 34.4032, 38.9020, 40.6494],\n",
       "        [35.3413, 32.8691, 37.6597, 39.0699, 41.6656, 34.0233, 38.4128, 40.1341]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "多头注意力就是增加了另外的q、k、v，从而可以从多个角度来捕捉token之间的相关性。",
   "id": "4605d9dfe0d35435"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-27T10:02:50.023027Z",
     "start_time": "2025-07-27T10:02:50.012644Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class GQAAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim: int, attn_dim: int, output_dim: int, num_heads: int, num_kv_groups: int):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.attn_dim = attn_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.num_kv_groups = num_kv_groups  # 新增\n",
    "        self.head_dim = attn_dim // num_heads\n",
    "\n",
    "        # QKV投影层：从输入维度映射到内部维度\n",
    "        self.q_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "\n",
    "        self.kv_attn_dim = self.head_dim * self.num_kv_groups\n",
    "        self.k_proj = nn.Linear(embed_dim, self.kv_attn_dim)\n",
    "        self.v_proj = nn.Linear(embed_dim, self.kv_attn_dim)\n",
    "\n",
    "        # 输出投影层：从内部维度映射到输出维度\n",
    "        self.out_proj = nn.Linear(self.attn_dim, self.output_dim)\n",
    "\n",
    "    def forward(self, x, kv_cache=None):\n",
    "        \"\"\"\n",
    "        输入: [batch_size, seq_len, embed_dim]\n",
    "        返回: [batch_size, seq_len, output_dim]\n",
    "        \"\"\"\n",
    "        batch_size, seq_len, embed_dim = x.shape\n",
    "\n",
    "        # 投影到QKV空间\n",
    "        q = self.q_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        k = self.k_proj(x)  # [batch_size, seq_len, kv_attn_dim]\n",
    "        v = self.v_proj(x)  # [batch_size, seq_len, kv_attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        # 分割多头 [batch_size, num_heads, seq_len, head_dim]\n",
    "        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        k = k.view(batch_size, seq_len, self.num_kv_groups, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        v = v.view(batch_size, seq_len, self.num_kv_groups, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        rep_factor = self.num_heads // self.num_kv_groups\n",
    "        k = k.repeat_interleave(rep_factor, dim=1)\n",
    "        v = v.repeat_interleave(rep_factor, dim=1)\n",
    "\n",
    "        if kv_cache is not None:\n",
    "            # 获取缓存的kv\n",
    "            k_cache, v_cache = kv_cache\n",
    "            # 拼接缓存的kv和当前的kv\n",
    "            k = torch.cat([k_cache, k], dim=2)\n",
    "            v = torch.cat([v_cache, v], dim=2)\n",
    "            # 更新缓存的kv\n",
    "            kv_cache = (k, v)\n",
    "        else:\n",
    "            kv_cache = (k, v)\n",
    "\n",
    "        # 计算注意力得分\n",
    "        # q   [batch_size, num_heads, seq_len, head_dim]\n",
    "        # k.T [batch_size, num_heads, head_dim, seq_len]\n",
    "        # q @ k.T 形状: [batch_size, num_heads, seq_len, seq_len]\n",
    "        # 1,4,4,2\n",
    "        # 1,4,2,4\n",
    "        attn_scores = torch.matmul(q, k.transpose(-2, -1))\n",
    "\n",
    "        # 缩放因子：防止乘积过大\n",
    "        d_k = k.size(-1)\n",
    "        attn_scores = attn_scores / torch.sqrt(torch.tensor(d_k))\n",
    "\n",
    "        # 计算注意力权重\n",
    "        attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "\n",
    "        # 计算注意力输出\n",
    "        # attn_weights [batch_size, num_heads, seq_len, seq_len]\n",
    "        # v            [batch_size, num_heads, seq_len, head_dim]\n",
    "        # [batch_size, num_heads, seq_len, head_dim]\n",
    "        attn_out = torch.matmul(attn_weights, v)\n",
    "\n",
    "        # 合并多头 [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        attn_out = attn_out.transpose(1, 2).reshape(batch_size, seq_len, self.attn_dim)\n",
    "\n",
    "        # 投影到输出空间\n",
    "        return self.out_proj(attn_out), kv_cache"
   ],
   "id": "403c5a8fcae45aae",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-27T10:02:54.662221Z",
     "start_time": "2025-07-27T10:02:54.654973Z"
    }
   },
   "cell_type": "code",
   "source": [
    "x = torch.rand(1, 4, 2)\n",
    "\n",
    "attn = GQAAttention(embed_dim=2, attn_dim=8, output_dim=8, num_heads=4, num_kv_groups=2)\n",
    "out, kv_cache = attn(x)\n",
    "\n",
    "print(x.shape)\n",
    "print(out.shape)\n",
    "\n",
    "x = torch.rand(1, 5, 2)\n",
    "out, kv_cache = attn(x[:, -1:, :], kv_cache)\n",
    "print(x.shape)\n",
    "print(out.shape)"
   ],
   "id": "6c229fa626c5ff1c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 2])\n",
      "torch.Size([1, 4, 8])\n",
      "torch.Size([1, 5, 2])\n",
      "torch.Size([1, 1, 8])\n"
     ]
    }
   ],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-27T10:04:00.607671Z",
     "start_time": "2025-07-27T10:04:00.600609Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "\n",
    "batch_size = 1\n",
    "seq_len = 1\n",
    "num_heads = 4\n",
    "head_dim = 1\n",
    "num_kv_headers = 2\n",
    "\n",
    "# 测试repeat函数效果\n",
    "q = torch.rand(batch_size, num_heads, seq_len, head_dim)\n",
    "print(q)\n",
    "\n",
    "k = torch.rand(batch_size, num_kv_headers, seq_len, head_dim)\n",
    "print(k)\n",
    "\n",
    "rep_factor = num_heads // num_kv_headers\n",
    "k_repeated = k.repeat_interleave(rep_factor, dim=1)\n",
    "print(k_repeated)\n",
    "\n",
    "attn_scores = torch.matmul(q, k_repeated.transpose(-2, -1))\n",
    "print(attn_scores.shape)"
   ],
   "id": "6d479aa88bc33377",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[0.4041]],\n",
      "\n",
      "         [[0.6990]],\n",
      "\n",
      "         [[0.2889]],\n",
      "\n",
      "         [[0.5308]]]])\n",
      "tensor([[[[0.4436]],\n",
      "\n",
      "         [[0.8200]]]])\n",
      "tensor([[[[0.4436]],\n",
      "\n",
      "         [[0.4436]],\n",
      "\n",
      "         [[0.8200]],\n",
      "\n",
      "         [[0.8200]]]])\n",
      "torch.Size([1, 4, 1, 1])\n"
     ]
    }
   ],
   "execution_count": 13
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
